﻿using System;
using System.IdentityModel.Tokens;
using SecurityTokenServiceNS;
using System.Security.Cryptography;
using System.ServiceModel.Security.Tokens;
using System.ServiceModel;

namespace ProofAddIns
{
	public class ProofModule : IProofProcessor
    {
        /// <summary>
        /// This method generates a Proof Token
        /// </summary>
        /// <param name="rst">RST</param>
        /// <param name="rstr">RSTR</param>
        /// <return>Proof Token</param>
        public SecurityKeyIdentifier GenerateProof(RST rst, RSTR rstr)
        {
            if (rst == null)
                throw new ArgumentNullException("rst");

            if (rstr == null)
                throw new ArgumentNullException("rstr");

			if (rst.ProofType == Constants.WSTrust.ProofType.PublicKey)
				return GenerateAsymmetricProof(rst);
			else if (rst.ProofType == Constants.WSTrust.ProofType.SymmetricKey)
				return GenerateSymmetricProof(rst, rstr);
			else
				return null;
        }

        private SecurityKeyIdentifier GenerateAsymmetricProof(RST rst)
        {
			return new SecurityKeyIdentifier(rst.UseKey.CreateKeyIdentifierClause<RsaKeyIdentifierClause>());
        }

        private SecurityKeyIdentifier GenerateSymmetricProof(RST rst, RSTR rstr)
        {
            byte[] sessionKey = GetSessionKey(rst, rstr);

            X509SecurityToken rpToken = (X509SecurityToken)StsConfiguration.Current.GetRpToken(rst.AppliesTo);
            byte[] wrappedKey = rpToken.SecurityKeys[0].EncryptKey(SecurityAlgorithms.RsaOaepKeyWrap, sessionKey);
            EncryptedKeyIdentifierClause clause = new EncryptedKeyIdentifierClause(wrappedKey, SecurityAlgorithms.RsaOaepKeyWrap,
                new SecurityKeyIdentifier(rpToken.CreateKeyIdentifierClause<X509ThumbprintKeyIdentifierClause>()));

            return new SecurityKeyIdentifier(clause);
        }

        private byte[] GetSessionKey(RST rst, RSTR rstr)
        {
            RNGCryptoServiceProvider random = new RNGCryptoServiceProvider();
            BinarySecretSecurityToken clientEntropy = rst.Entropy as BinarySecretSecurityToken;
            int keySize = rst.KeySize != 0 ? rst.KeySize : 256;

            byte[] sessionKey = null;
            byte[] issuerEntropy = new byte[keySize / 8];
            random.GetNonZeroBytes(issuerEntropy);

            if (clientEntropy != null)  // Combined entropy
            {
                sessionKey = ComputeCombinedKey(clientEntropy.GetKeyBytes(), issuerEntropy, keySize);
                rstr.Entropy = new BinarySecretSecurityToken(issuerEntropy);
                rstr.ComputedKey = true;
            }
            else                        // Issuer-only entropy
            {
                sessionKey = issuerEntropy;
                rstr.RequestedProofToken = new BinarySecretSecurityToken(sessionKey);
            }

            rstr.KeySize = keySize;
            return sessionKey;
        }

        private byte[] ComputeCombinedKey(byte[] requestorEntropy, byte[] issuerEntropy, int keySize)
        {
            if (keySize < 64 || keySize > 4096)
                throw new ArgumentOutOfRangeException("keySize");

            KeyedHashAlgorithm kha = new HMACSHA1(requestorEntropy, true);

            byte[] key = new byte[keySize / 8];                 // Final key
            byte[] a = issuerEntropy;                           // A(0)
            byte[] b = new byte[kha.HashSize / 8 + a.Length];   // Buffer for A(i) + seed

            for (int i = 0; i < key.Length; )
            {
                // Calculate A(i+1).                
                kha.Initialize();
                a = kha.ComputeHash(a);

                // Calculate A(i) + seed
                a.CopyTo(b, 0);
                issuerEntropy.CopyTo(b, a.Length);
                kha.Initialize();
                byte[] result = kha.ComputeHash(b);

                for (int j = 0; j < result.Length; j++)
                {
                    if (i < key.Length)
                        key[i++] = result[j];
                    else
                        break;
                }
            }
            return key;
        }
    }
}
